"""
跨网络数据传输服务器
支持多个不同网络环境下的Android设备实时数据同步
"""

import asyncio
import json
import logging
import os
from datetime import datetime
from typing import Dict, Set, List
import websockets
import socket

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class CrossNetworkDataServer:
    def __init__(self):
        # 连接管理
        self.clients: Set[websockets.WebSocketServerProtocol] = set()
        self.user_connections: Dict[str, websockets.WebSocketServerProtocol] = {}
        self.department_connections: Dict[str, Set[websockets.WebSocketServerProtocol]] = {}
        
        # 数据存储文件路径
        self.data_dir = "server_data"
        self.patients_file = os.path.join(self.data_dir, "patients.json")
        self.users_file = os.path.join(self.data_dir, "users.json")
        self.notifications_file = os.path.join(self.data_dir, "notifications.json")
        
        # 确保数据目录存在
        os.makedirs(self.data_dir, exist_ok=True)
        
        # 数据存储（支持持久化）
        self.patients_data: List[Dict] = self.load_data(self.patients_file)
        self.users_data: List[Dict] = self.load_data(self.users_file)
        self.notifications_data: List[Dict] = self.load_data(self.notifications_file)
        
        # 网络环境标识
        self.client_networks: Dict[str, str] = {}
        
        logger.info(f"已加载数据: {len(self.patients_data)}个患者, {len(self.users_data)}个用户, {len(self.notifications_data)}个通知")
        
    def load_data(self, file_path: str) -> List[Dict]:
        """从文件加载数据"""
        try:
            if os.path.exists(file_path):
                with open(file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
        except Exception as e:
            logger.error(f"加载数据文件失败 {file_path}: {e}")
        return []
    
    def save_data(self, data: List[Dict], file_path: str):
        """保存数据到文件"""
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            logger.error(f"保存数据文件失败 {file_path}: {e}")
    
    def save_all_data(self):
        """保存所有数据到文件"""
        self.save_data(self.patients_data, self.patients_file)
        self.save_data(self.users_data, self.users_file)
        self.save_data(self.notifications_data, self.notifications_file)
        logger.info("所有数据已保存到文件")
        
    async def register_client(self, websocket: websockets.WebSocketServerProtocol):
        """注册新客户端"""
        self.clients.add(websocket)
        client_address = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        
        # 识别网络环境
        network_type = self.identify_network_type(websocket.remote_address[0])
        self.client_networks[client_address] = network_type
        
        logger.info(f"新客户端连接: {client_address} ({network_type})，当前连接数: {len(self.clients)}")
        
        # 发送欢迎消息
        welcome_message = {
            'type': 'SERVER_INFO',
            'data': {
                'message': '成功连接到跨网络数据传输服务器',
                'server_time': datetime.now().isoformat(),
                'network_type': network_type,
                'total_clients': len(self.clients)
            }
        }
        await websocket.send(json.dumps(welcome_message))
        
    def identify_network_type(self, ip_address: str) -> str:
        """识别网络类型"""
        if ip_address.startswith('192.168.'):
            return 'Local WiFi'
        elif ip_address.startswith('10.'):
            return 'Private Network'
        elif ip_address.startswith('172.'):
            return 'Corporate Network'
        else:
            return 'Public Network'
    
    async def unregister_client(self, websocket: websockets.WebSocketServerProtocol):
        """注销客户端"""
        self.clients.discard(websocket)
        client_address = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        
        # 从映射中移除
        if client_address in self.client_networks:
            del self.client_networks[client_address]
            
        # 从用户连接映射中移除
        user_id = None
        for uid, conn in list(self.user_connections.items()):
            if conn == websocket:
                user_id = uid
                break
        
        if user_id:
            del self.user_connections[user_id]
            
        # 从部门连接映射中移除
        for dept, connections in list(self.department_connections.items()):
            connections.discard(websocket)
            if not connections:
                del self.department_connections[dept]
                
        logger.info(f"客户端断开连接: {client_address}，当前连接数: {len(self.clients)}")
        
    async def handle_message(self, websocket: websockets.WebSocketServerProtocol, message):
        """处理接收到的消息"""
        try:
            # 处理消息类型转换，支持 bytes 和 str 类型
            if isinstance(message, bytes):
                message_str = message.decode('utf-8')
            else:
                message_str = str(message)
            
            data = json.loads(message_str)
            message_type = data.get('type')
            
            logger.info(f"收到消息类型: {message_type}")
            
            # 路由消息到对应处理器
            handlers = {
                'AUTH': self.handle_auth,
                'PATIENT_REPORT': self.handle_patient_report,
                'NOTIFICATION': self.handle_notification,
                'HEARTBEAT': self.handle_heartbeat,
                'DATA_SYNC': self.handle_data_sync,
                'USER_REGISTER': self.handle_user_register
            }
            
            handler = handlers.get(message_type, self.handle_unknown_message)
            await handler(websocket, data)
                
        except json.JSONDecodeError as e:
            await self.send_error(websocket, f"消息格式错误: {e}")
        except Exception as e:
            logger.error(f"处理消息时出错: {e}")
            await self.send_error(websocket, f"处理消息时出错: {str(e)}")
    
    async def send_error(self, websocket: websockets.WebSocketServerProtocol, error_msg: str):
        """发送错误消息"""
        error_response = {
            'type': 'ERROR',
            'data': {
                'message': error_msg,
                'timestamp': datetime.now().isoformat()
            }
        }
        await websocket.send(json.dumps(error_response))
    
    async def handle_unknown_message(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理未知消息类型"""
        message_type = data.get('type')
        logger.warning(f"未知消息类型: {message_type}")
        await self.send_error(websocket, f'未知消息类型: {message_type}')
            
    async def handle_auth(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理认证消息"""
        auth_data = data.get('data', {})
        user_id = str(auth_data.get('userId', ''))
        username = auth_data.get('username', '')
        user_type = auth_data.get('userType', '')
        department = auth_data.get('department', '')
        
        # 注册用户连接
        self.user_connections[user_id] = websocket
        
        # 注册部门连接
        if department:
            if department not in self.department_connections:
                self.department_connections[department] = set()
            self.department_connections[department].add(websocket)
            
        logger.info(f"用户认证成功: {username} ({user_type}) - {department}")
        
        # 发送认证确认和当前数据摘要
        response = {
            'type': 'AUTH_SUCCESS',
            'data': {
                'userId': user_id,
                'timestamp': datetime.now().isoformat(),
                'data_summary': {
                    'patients_count': len(self.patients_data),
                    'users_count': len(self.users_data),
                    'notifications_count': len(self.notifications_data)
                }
            }
        }
        await websocket.send(json.dumps(response))
        
    async def handle_patient_report(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理患者上报消息"""
        report_data = data.get('data', {})
        patient_id = report_data.get('patientId', '')
        patient_name = report_data.get('patientName', '')
        reporter_name = report_data.get('reporterName', '')
        department = report_data.get('department', '')
        timestamp = report_data.get('timestamp', datetime.now().timestamp() * 1000)
        
        # 存储患者数据
        patient_record = {
            'id': patient_id,
            'name': patient_name,
            'reporter': reporter_name,
            'department': department,
            'timestamp': timestamp,
            'report_time': datetime.now().isoformat()
        }
        
        # 检查是否已存在，避免重复
        existing_patient = next((p for p in self.patients_data if p['id'] == patient_id), None)
        if not existing_patient:
            self.patients_data.append(patient_record)
            # 自动保存数据
            self.save_data(self.patients_data, self.patients_file)
            logger.info(f"新患者上报已保存: {patient_name} - {reporter_name} ({department})")
        else:
            logger.info(f"患者记录已存在: {patient_name}")
        
        # 广播给所有其他客户端（实现跨网络同步）
        notification = {
            'type': 'PATIENT_REPORT_BROADCAST',
            'data': patient_record
        }
        
        await self.broadcast_to_others(websocket, notification)
        
        # 发送确认给发送方
        confirmation = {
            'type': 'PATIENT_REPORT_SUCCESS',
            'data': {
                'patientId': patient_id,
                'timestamp': datetime.now().isoformat(),
                'synced_to_clients': len(self.clients) - 1
            }
        }
        await websocket.send(json.dumps(confirmation))
        
    async def handle_user_register(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理用户注册消息"""
        register_data = data.get('data', {})
        user_id = register_data.get('userId', '')
        username = register_data.get('username', '')
        real_name = register_data.get('realName', '')
        department = register_data.get('department', '')
        user_type = register_data.get('userType', '')
        phone = register_data.get('phone', '')
        
        # 存储用户数据
        user_record = {
            'id': user_id,
            'username': username,
            'realName': real_name,
            'department': department,
            'userType': user_type,
            'phone': phone,
            'register_time': datetime.now().isoformat()
        }
        
        self.users_data.append(user_record)
        # 自动保存数据
        self.save_data(self.users_data, self.users_file)
        logger.info(f"新用户注册: {username} ({real_name}) - {department}")
        
        # 广播给管理员
        notification = {
            'type': 'USER_REGISTER_BROADCAST',
            'data': user_record
        }
        await self.broadcast_to_admins(notification)
        
    async def handle_notification(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理通知消息"""
        notification_data = data.get('data', {})
        title = notification_data.get('title', '')
        content = notification_data.get('content', '')
        target_user_id = notification_data.get('targetUserId')
        target_department = notification_data.get('targetDepartment')
        
        # 存储通知
        notification_record = {
            'title': title,
            'content': content,
            'targetUserId': target_user_id,
            'targetDepartment': target_department,
            'timestamp': datetime.now().isoformat()
        }
        self.notifications_data.append(notification_record)
        # 自动保存数据
        self.save_data(self.notifications_data, self.notifications_file)
        
        logger.info(f"收到通知: {title} - {content}")
        
        notification = {
            'type': 'NOTIFICATION_BROADCAST',
            'data': notification_record
        }
        
        # 路由通知
        if target_user_id:
            await self.send_to_user(str(target_user_id), notification)
        elif target_department:
            await self.broadcast_to_department(target_department, notification)
        else:
            await self.broadcast_to_all(notification)
            
    async def handle_heartbeat(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理心跳消息"""
        response = {
            'type': 'HEARTBEAT_RESPONSE',
            'data': {
                'timestamp': datetime.now().isoformat(),
                'server_status': 'healthy',
                'connected_clients': len(self.clients)
            }
        }
        await websocket.send(json.dumps(response))
    
    async def handle_data_sync(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理数据同步请求"""
        sync_data = data.get('data', {})
        sync_type = sync_data.get('type', '')
        
        if sync_type == 'request':
            # 发送全部数据给请求方
            full_data = {
                'type': 'DATA_SYNC_RESPONSE',
                'data': {
                    'patients': self.patients_data,
                    'users': self.users_data,
                    'notifications': self.notifications_data,
                    'timestamp': datetime.now().isoformat()
                }
            }
            await websocket.send(json.dumps(full_data))
            logger.info(f"发送完整数据同步: {len(self.patients_data)}个患者, {len(self.users_data)}个用户")
            
        elif sync_type == 'patients':
            # 接收患者数据同步
            patients_data = json.loads(sync_data.get('data', '[]'))
            for patient in patients_data:
                existing = next((p for p in self.patients_data if p.get('id') == patient.get('id')), None)
                if not existing:
                    self.patients_data.append(patient)
            logger.info(f"同步了 {len(patients_data)} 条患者数据")
            
        elif sync_type == 'users':
            # 接收用户数据同步
            users_data = json.loads(sync_data.get('data', '[]'))
            for user in users_data:
                existing = next((u for u in self.users_data if u.get('id') == user.get('id')), None)
                if not existing:
                    self.users_data.append(user)
            logger.info(f"同步了 {len(users_data)} 条用户数据")
    
    async def send_to_user(self, user_id: str, message: dict):
        """发送消息给指定用户"""
        connection = self.user_connections.get(user_id)
        if connection and connection in self.clients:
            try:
                await connection.send(json.dumps(message))
            except:
                # 连接已断开，清理
                if user_id in self.user_connections:
                    del self.user_connections[user_id]
    
    async def broadcast_to_all(self, message: dict):
        """广播给所有客户端"""
        if not self.clients:
            return
            
        message_str = json.dumps(message)
        disconnected = []
        
        for client in self.clients:
            try:
                await client.send(message_str)
            except:
                disconnected.append(client)
        
        # 清理断开的连接
        for client in disconnected:
            self.clients.discard(client)
    
    async def broadcast_to_others(self, sender: websockets.WebSocketServerProtocol, message: dict):
        """广播给除发送者外的所有客户端"""
        others = self.clients - {sender}
        if not others:
            return
            
        message_str = json.dumps(message)
        disconnected = []
        
        for client in others:
            try:
                await client.send(message_str)
            except:
                disconnected.append(client)
        
        # 清理断开的连接
        for client in disconnected:
            self.clients.discard(client)
            
    async def broadcast_to_admins(self, message: dict):
        """广播给管理员"""
        admin_connections = set()
        for user_id, connection in self.user_connections.items():
            if connection in self.clients:
                admin_connections.add(connection)
            
        if admin_connections:
            message_str = json.dumps(message)
            for client in admin_connections:
                try:
                    await client.send(message_str)
                except:
                    pass
            
    async def broadcast_to_department(self, department: str, message: dict):
        """广播给指定部门"""
        if department not in self.department_connections:
            return
            
        connections = self.department_connections[department] & self.clients
        if connections:
            message_str = json.dumps(message)
            for client in connections:
                try:
                    await client.send(message_str)
                except:
                    pass
            
    async def handler(self, websocket: websockets.WebSocketServerProtocol, path: str):
        """WebSocket连接处理器"""
        await self.register_client(websocket)
        
        try:
            async for message in websocket:
                await self.handle_message(websocket, message)
        except websockets.exceptions.ConnectionClosed:
            logger.info("连接正常关闭")
        except Exception as e:
            logger.error(f"连接处理出错: {e}")
        finally:
            await self.unregister_client(websocket)

def get_local_ip():
    """获取本机IP地址"""
    try:
        # 连接到一个远程地址来获取本机IP
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))
            return s.getsockname()[0]
    except:
        return "127.0.0.1"

async def main():
    """主函数"""
    server = CrossNetworkDataServer()
    
    # 启动WebSocket服务器
    host = "0.0.0.0"  # 监听所有接口
    port = 8085  # 使用8085端口
    local_ip = get_local_ip()
    
    logger.info("=" * 60)
    logger.info("跨网络数据传输服务器启动中...")
    logger.info(f"本机IP地址: {local_ip}")
    logger.info(f"监听地址: ws://{host}:{port}")
    logger.info(f"外网访问: ws://8.140.59.178:{port}")
    logger.info("支持多网络环境设备实时数据同步")
    logger.info("=" * 60)
    
    async with websockets.serve(server.handler, host, port, ping_interval=30, ping_timeout=10):
        logger.info("WebSocket服务器启动成功，等待客户端连接...")
        await asyncio.Future()  # 保持服务器运行

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("服务器停止")
        # 保存数据在关闭前
        try:
            server = CrossNetworkDataServer()
            server.save_all_data()
            logger.info("所有数据已保存")
        except:
            pass
    except Exception as e:
        logger.error(f"服务器运行出错: {e}")